#!/usr/bin/python

from datetime import timedelta
import logging

from hips import datastore
from hips.scheduling import constraint
from hips.scheduling import doctor
from hips.scheduling import problem
from hips.scheduling import schedule


def CreateHipsSchedulingProblem(ds_schedule, ds_schedule_block_list,
    ds_doctor_constraint_list):
  """Given a datastore schedule and constraints, creates a HIPS problem.

  Does not execute the problem that it returns.

  Args:
    ds_schedule: Datastore schedule on which to base problem.
    ds_schedule_block_list: Schedule blocks for above schedule.
    ds_doctor_constraint_list: Constraints for problem.

  Returns:
    SchedulingProblem instance.
  """
  hips_block_list = []
  for ds_block in ds_schedule_block_list:
    hips_block = schedule.ScheduleBlock(
        ds_block.start_time, timedelta(minutes=ds_block.duration_minutes))
    hips_block_list.append(hips_block)
  hips_schedule = schedule.Schedule(hips_block_list)
  p = problem.SchedulingProblem(hips_schedule)
  for i, d in enumerate(ds_doctor_constraint_list):
    logging.info("Creating doctor %d: %s" % (i, d.doctor.name))
    logging.info("Days off: %s" % str(d.days_off))
    hips_doctor = doctor.Doctor(
        d.doctor.name,
        days_off=d.days_off,
        min_shifts=d.min_shifts,
        max_shifts=d.max_shifts,
        max_night_shifts=d.max_night_shifts,
        max_consecutive_days=d.max_consecutive_days)
    p.AddDoctor(hips_doctor)
  return p


def CreateScheduleSolutionModels(ds_schedule, ds_schedule_block_list,
    hips_solution, doctor_name_to_ds_key_map):
  """Given a HIPS solution instance, creates datastore models.

  The datastore forces you to save models before you can refer to them
  by key, so this function saves all created models before returning.

  Args:
    ds_schedule: Schedule that the solution is based upon.
    ds_schedule_block_list: Schedule blocks for above schedule.
    hips_solution: HIPS solution to format.
    doctor_name_to_ds_key_map: Map of <name, key> to join doctors on.

  Returns:
    A list of created models: [ScheduleSolution, ScheduleSolutionBlock*]
  """
  solution_map = hips_solution.FormatSolutionAsMap()
  # TODO Properly communicate stats, use FormatSolutionAsString for now
  ds_solution = datastore.ScheduleSolution(
      schedule=ds_schedule,
      message=hips_solution.FormatSolutionAsString())
  ds_solution.put()
  model_list = [ds_solution]
  solution_list = solution_map.get("solution")
  for single_shift_solution_list, ds_block in \
      zip(solution_list, ds_schedule_block_list):
    for hips_doctor in single_shift_solution_list:
      ds_doctor = doctor_name_to_ds_key_map.get(hips_doctor.name)
      ds_solution_block = datastore.ScheduleSolutionBlock(
          solution=ds_solution,
          block=ds_block,
          doctor=ds_doctor)
      ds_solution_block.put()
      model_list.append(ds_solution_block)
  return model_list


def CreateScheduleModels(hips_schedule, ds_schedule_model, ds_doctor_list):
  """Given a HIPS schedule instance, creates datastore models.

  The datastore forces you to save models before you can refer to them
  by key, so this function saves all created models before returning.

  Args:
    hips_schedule:
    ds_schedule_model:
    ds_doctor_list:

  Returns:
    A list of created models: [Schedule, ScheduleBlock*, DoctorConstraint*]
  """
  ds_schedule_model.put()
  model_list = [ds_schedule_model]
  for hips_schedule_block in hips_schedule.ScheduleBlocks():
    ds_schedule_block_model = datastore.ScheduleBlock(
        schedule=ds_schedule_model,
        start_time=hips_schedule_block.start_time,
        duration_minutes=hips_schedule_block.duration.seconds // 60)
    ds_schedule_block_model.put()
    model_list.append(ds_schedule_block_model)
  constraint_model_list = []
  for ds_doctor_model in ds_doctor_list:
    ds_constraint_model = datastore.DoctorConstraint(
        doctor=ds_doctor_model,
        schedule=ds_schedule_model,
        days_off=[],
        min_shifts=ds_doctor_model.min_shifts,
        max_shifts=ds_doctor_model.max_shifts,
        max_night_shifts=ds_doctor_model.max_night_shifts,
        max_consecutive_days=ds_doctor_model.max_consecutive_days)
    ds_constraint_model.put()
    model_list.append(ds_constraint_model)
  return model_list
